/*
 * Copyright (c) 2020-2099 the original author or authors. All rights reserve.
 */
package org.tianyun.cloud.security.token.store;

import org.tianyun.cloud.security.token.common.AccessToken;

import java.security.Principal;
import java.util.Date;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.Delayed;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Implementation of token services that stores tokens in memory.
 *
 * @auther ebert_chan
 */
public class InMemoryTokenStore implements TokenStore, AutoCloseable {

    private static final int DEFAULT_FLUSH_INTERVAL = 1000;

    private final ConcurrentHashMap<String, AccessToken> accessTokenStore = new ConcurrentHashMap<String, AccessToken>();

    private final DelayQueue<TokenExpiry> expiryQueue = new DelayQueue<TokenExpiry>();

    private final ConcurrentHashMap<String, TokenExpiry> expiryMap = new ConcurrentHashMap<String, TokenExpiry>();

    private int flushInterval = DEFAULT_FLUSH_INTERVAL;

    private AtomicInteger flushCounter = new AtomicInteger(0);

    /**
     * The number of tokens to store before flushing expired tokens. Defaults to 1000.
     *
     * @param flushInterval the interval to set
     */
    public void setFlushInterval(int flushInterval) {
        this.flushInterval = flushInterval;
    }

    /**
     * The interval (count of token inserts) between flushing expired tokens.
     *
     * @return the flushInterval the flush interval
     */
    public int getFlushInterval() {
        return flushInterval;
    }

    /*
     * @see org.tianyun.cloud.security.token.store.TokenStore#getAccessToken(java.security.Principal)
     */
    @Override
    public AccessToken getAccessToken(Principal principal) {
        AccessToken accessToken = this.accessTokenStore.get(principal.getName());
        if (accessToken != null && accessToken.isExpired()) {
            this.removeAccessToken(principal);
            accessToken = null;
        }
        return accessToken;
    }

    /*
     * @see org.tianyun.cloud.security.token.store.TokenStore#storeAccessToken(java.security.Principal,
     * org.tianyun.cloud.security.token.common.AccessToken)
     */
    @Override
    public void storeAccessToken(Principal principal, AccessToken accessToken) {
        if (this.flushCounter.incrementAndGet() >= this.flushInterval) {
            flush();
            this.flushCounter.set(0);
        }

        this.accessTokenStore.put(principal.getName(), accessToken);

        if (accessToken.getExpiration() != null) {
            TokenExpiry expiry = new TokenExpiry(principal.getName(), accessToken.getExpiration());
            // Remove existing expiry for this token if present
            expiryQueue.remove(expiryMap.put(principal.getName(), expiry));
            this.expiryQueue.put(expiry);
        }

    }

    private void flush() {
        TokenExpiry expiry = expiryQueue.poll();
        while (expiry != null) {
            this.accessTokenStore.remove(expiry.getValue());
            expiry = expiryQueue.poll();
        }
    }

    /*
     * @see java.lang.AutoCloseable#close()
     */
    @Override
    public void close() throws Exception {
        accessTokenStore.clear();
        expiryMap.clear();
        expiryQueue.clear();
    }

    /*
     * @see org.tianyun.cloud.security.token.store.TokenStore#removeAccessToken(java.security.Principal)
     */
    @Override
    public void removeAccessToken(Principal principal) {
        this.accessTokenStore.remove(principal.getName());
    }

    private static class TokenExpiry implements Delayed {

        private final long expiry;

        private final String value;

        public TokenExpiry(String value, Date date) {
            this.value = value;
            this.expiry = date.getTime();
        }

        public int compareTo(Delayed other) {
            if (this == other) {
                return 0;
            }
            long diff = getDelay(TimeUnit.MILLISECONDS) - other.getDelay(TimeUnit.MILLISECONDS);
            return (diff == 0 ? 0 : ((diff < 0) ? -1 : 1));
        }

        public long getDelay(TimeUnit unit) {
            return expiry - System.currentTimeMillis();
        }

        public String getValue() {
            return value;
        }

    }

}
